use std::collections::HashSet;

use sea_orm::sea_query::Expr;
use sea_orm::{
    ActiveModelTrait, ActiveValue, ColumnTrait, ConnectionTrait, DatabaseConnection,
    DatabaseTransaction, EntityTrait, PaginatorTrait, QueryFilter, QuerySelect, Set, Statement,
};
use serde_json::Value;
use uuid::Uuid;

use crate::entities::policy::Policy;
use crate::entities::policy_db_model::{ActiveModel, Column, Entity as PolicyEntity, Model};
use crate::entities::signature_policy::SignaturePolicy;
use crate::policy_error::policy_error::PolicyError;
use key_management::key_manager::error::KeyManagerError;

pub struct PolicyRepository;

impl PolicyRepository {
    /// Gets the version of a policy by its ID.
    ///
    /// # Arguments
    /// * `db` - The database transaction
    /// * `policy_id` - The ID of the policy to query
    ///
    /// # Returns
    /// * `Result<i32, PolicyError>` - The policy version if found, otherwise an error
    pub async fn get_policy_version(
        db: &DatabaseTransaction,
        policy_id: String,
    ) -> Result<i32, PolicyError> {
        let result = PolicyEntity::find()
            .filter(Column::PolicyId.eq(policy_id))
            .column(Column::PolicyVersion)
            .into_model::<crate::entities::policy_db_model::Model>()
            .one(db)
            .await
            .map_err(|e| PolicyError::DatabaseOperationError(e.to_string()))?;
        result
            .map(|m| m.policy_version)
            .ok_or_else(|| PolicyError::PolicyNotFoundError("Policy not found".to_string()))
    }

    /// Adds a new policy to the database.
    ///
    /// # Arguments
    /// * `db` - The database transaction
    /// * `policy` - The policy model to be added
    ///
    /// # Returns
    /// * `Result<(), PolicyError>` - Ok if successful, otherwise a PolicyError
    pub async fn add_policy(
        db: &DatabaseTransaction,
        policy: ActiveModel,
    ) -> Result<(), PolicyError> {
        let policy_id = policy.policy_id.unwrap();
        let policy_name = policy.policy_name.unwrap();
        let policy_description = policy.policy_description.unwrap();
        let policy_content = policy.policy_content.unwrap();
        let is_default = policy.is_default.unwrap();
        let policy_version = policy.policy_version.unwrap();
        let create_time = policy.create_time.unwrap();
        let update_time = policy.update_time.unwrap();
        let user_id = policy.user_id.unwrap();
        let attester_type = policy.attester_type.unwrap();
        let signature = policy.signature.unwrap();
        let key_version = policy.key_version.unwrap();
        let product_name = policy.product_name.unwrap();
        let product_type = policy.product_type.unwrap();
        let board_type = policy.board_type.unwrap();

        let sql = r#"
            INSERT INTO policy_information (
                policy_id, policy_name, policy_description, policy_content,
                is_default, policy_version, create_time, update_time,
                user_id, attester_type, signature, valid_code, key_version,
                product_name, product_type, board_type
            ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
        "#;

        db.query_one(Statement::from_sql_and_values(
            db.get_database_backend(),
            sql,
            vec![
                policy_id.into(),
                policy_name.into(),
                policy_description.into(),
                policy_content.into(),
                is_default.into(),
                policy_version.into(),
                create_time.into(),
                update_time.into(),
                user_id.into(),
                attester_type.into(),
                signature.into(),
                0.into(),
                key_version.into(),
                product_name.into(),
                product_type.into(),
                board_type.into(),
            ],
        ))
        .await
        .map_err(|e| PolicyError::DatabaseOperationError(e.to_string()))?;

        Ok(())
    }

    /// Updates an existing policy in the database.
    ///
    /// # Arguments
    /// * `db` - The database transaction
    /// * `policy` - The policy model with updated values
    ///
    /// # Returns
    /// * `Result<Policy, PolicyError>` - The updated policy if successful, otherwise a PolicyError
    pub async fn update_policy(
        db: &DatabaseTransaction,
        policy: ActiveModel,
    ) -> Result<Policy, PolicyError> {
        let result = policy
            .update(db)
            .await
            .map_err(|e| PolicyError::DatabaseOperationError(e.to_string()))?;
        Ok(result.into())
    }

    /// Checks if a policy exists based on user ID and policy ID.
    ///
    /// # Arguments
    /// * `headers` - The HTTP request headers containing user ID
    /// * `db` - The database transaction
    /// * `request_body` - The request body containing policy ID
    ///
    /// # Returns
    /// * `Result<Option<Model>, PolicyError>` - The policy if found, None if not found, or PolicyError on failure
    pub async fn check_policy_exists(
        headers: &actix_web::http::header::HeaderMap,
        db: &DatabaseTransaction,
        request_body: &Value,
    ) -> Result<Option<Model>, PolicyError> {
        let user_id = headers
            .get("user_id")
            .and_then(|h| h.to_str().ok())
            .unwrap_or("system")
            .to_string();
        let id = request_body["id"].as_str().unwrap().to_string();
        let policy = PolicyEntity::find()
            .filter(Column::UserId.eq(user_id).and(Column::PolicyId.eq(id)))
            .one(db)
            .await
            .map_err(|e| PolicyError::DatabaseOperationError(e.to_string()))?;

        Ok(policy)
    }

    /// Deletes multiple policies by their IDs for a specific user.
    ///
    /// # Arguments
    /// * `db` - The database transaction
    /// * `policy_ids` - List of policy IDs to delete
    /// * `user_id` - The ID of the user who owns the policies
    ///
    /// # Returns
    /// * `Result<(), PolicyError>` - Ok if successful, PolicyError if any policy doesn't exist or on database error
    pub async fn delete_policies(
        db: &DatabaseTransaction,
        policy_ids: Vec<String>,
        user_id: String,
    ) -> Result<(), PolicyError> {
        let mut not_exist_policy_ids = Vec::new();
        let existing_policies = PolicyEntity::find()
            .filter(
                Column::PolicyId
                    .is_in(policy_ids.clone())
                    .and(Column::UserId.eq(user_id.clone())),
            )
            .all(db)
            .await
            .map_err(|e| PolicyError::DatabaseOperationError(e.to_string()))?;

        let existing_ids: HashSet<String> = existing_policies
            .iter()
            .map(|p| p.policy_id.clone())
            .collect();

        for uuid in policy_ids {
            if !existing_ids.contains(&uuid) {
                not_exist_policy_ids.push(uuid.clone());
            }
        }
        if !not_exist_policy_ids.is_empty() {
            return Err(PolicyError::PolicyNotFoundError(
                format!("Can't find policies with IDs: {}", not_exist_policy_ids.join(", "))
            ));
        }

        if !existing_ids.is_empty() {
            PolicyEntity::delete_many()
                .filter(Column::PolicyId.is_in(existing_ids))
                .exec(db)
                .await
                .map_err(|e| PolicyError::DatabaseOperationError(e.to_string()))?;
        }

        Ok(())
    }

    /// Deletes all policies of a specific attester type for a user.
    ///
    /// # Arguments
    /// * `db` - The database transaction
    /// * `attester_type` - The type of attester whose policies to delete
    /// * `user_id` - The ID of the user who owns the policies
    ///
    /// # Returns
    /// * `Result<(), PolicyError>` - Ok if successful, otherwise a PolicyError
    pub async fn delete_policies_by_type(
        db: &DatabaseTransaction,
        attester_type: String,
        user_id: String,
    ) -> Result<(), PolicyError> {
        let result = PolicyEntity::delete_many()
            .filter(Column::UserId.eq(user_id).and(Expr::cust(&format!(
                "JSON_CONTAINS(attester_type, '\"{}\"')",
                attester_type
            ))))
            .exec(db)
            .await
            .map_err(|e| PolicyError::DatabaseOperationError(e.to_string()))?;
        
        if result.rows_affected == 0 {
            return Err(PolicyError::PolicyNotFoundError(
                format!("No policies found for attester type: {}", attester_type)
            ));
        }
        Ok(())
    }

    /// Deletes all policies for a specific user.
    ///
    /// # Arguments
    /// * `db` - The database transaction
    /// * `user_id` - The ID of the user whose policies to delete
    ///
    /// # Returns
    /// * `Result<(), PolicyError>` - Ok if successful, otherwise a PolicyError
    pub async fn delete_all_policies(
        db: &DatabaseTransaction,
        user_id: String,
    ) -> Result<(), PolicyError> {
        let result = PolicyEntity::delete_many()
            .filter(Column::UserId.eq(user_id.clone()))
            .exec(db)
            .await
            .map_err(|e| PolicyError::DatabaseOperationError(e.to_string()))?;
        
        if result.rows_affected == 0 {
            return Err(PolicyError::PolicyNotFoundError(
                format!("No policies found for user: {}", user_id)
            ));
        }
        Ok(())
    }

    /// Retrieves all default policies for a specific attester type.
    ///
    /// # Arguments
    /// * `db` - The database connection
    /// * `attester_type` - The type of attester to filter policies
    ///
    /// # Returns
    /// * `Result<Vec<SignaturePolicy>, PolicyError>` - List of default policies if successful, otherwise a PolicyError
    pub async fn get_default_policies_by_type(
        db: &DatabaseConnection,
        attester_type: String,
    ) -> Result<Vec<SignaturePolicy>, PolicyError> {
        let policy_models = PolicyEntity::find()
            .filter(
                Column::IsDefault
                    .eq(true)
                    .and(Column::ValidCode.eq(0))
                    .and(Expr::cust(&format!(
                        "JSON_CONTAINS(attester_type, '\"{}\"')",
                        attester_type
                    ))),
            )
            .into_model::<Model>()
            .all(db)
            .await
            .map_err(|e| PolicyError::DatabaseOperationError(e.to_string()))?;

        let signature_policies: Vec<SignaturePolicy> = policy_models
            .into_iter()
            .map(|model| model.into())
            .collect::<Vec<_>>();

        Ok(signature_policies)
    }

    /// Retrieves multiple policies by their IDs for a specific user.
    ///
    /// # Arguments
    /// * `db` - The database transaction
    /// * `uuids` - List of policy IDs to retrieve
    /// * `user_id` - The ID of the user who owns the policies
    ///
    /// # Returns
    /// * `Result<Vec<SignaturePolicy>, PolicyError>` - List of policies if successful, otherwise a PolicyError
    pub async fn get_policies_by_ids(
        db: &DatabaseTransaction,
        uuids: Vec<String>,
        user_id: String,
    ) -> Result<Vec<SignaturePolicy>, PolicyError> {
        let policy_models = PolicyEntity::find()
            .filter(
                Column::PolicyId
                    .is_in(uuids)
                    .and(Column::UserId.eq(user_id)),
            )
            .into_model::<Model>()
            .all(db)
            .await
            .map_err(|e| PolicyError::DatabaseOperationError(e.to_string()))?;

        let policies: Vec<SignaturePolicy> = policy_models
            .into_iter()
            .map(|model| SignaturePolicy::from(model))
            .collect();

        Ok(policies)
    }

    /// Retrieves all policies of a specific attester type for a user.
    ///
    /// # Arguments
    /// * `db` - The database transaction
    /// * `attester_type` - The type of attester to filter policies
    /// * `user_id` - The ID of the user who owns the policies
    ///
    /// # Returns
    /// * `Result<Vec<SignaturePolicy>, PolicyError>` - List of policies if successful, otherwise a PolicyError
    pub async fn get_policies_by_type(
        db: &DatabaseTransaction,
        attester_type: String,
        user_id: String,
    ) -> Result<Vec<SignaturePolicy>, PolicyError> {
        let policy_models = PolicyEntity::find()
            .filter(Column::UserId.eq(user_id).and(Expr::cust(&format!(
                "JSON_CONTAINS(attester_type, '\"{}\"')",
                attester_type
            ))))
            .into_model::<Model>()
            .all(db)
            .await
            .map_err(|e| PolicyError::DatabaseOperationError(e.to_string()))?;
        let policies: Vec<SignaturePolicy> = policy_models
            .into_iter()
            .map(|model| SignaturePolicy::from(model))
            .collect();

        Ok(policies)
    }

    /// Retrieves all policies for a specific user.
    ///
    /// # Arguments
    /// * `db` - The database transaction
    /// * `user_id` - The ID of the user whose policies to retrieve
    ///
    /// # Returns
    /// * `Result<Vec<SignaturePolicy>, PolicyError>` - List of all policies if successful, otherwise a PolicyError
    pub async fn get_all_policies(
        db: &DatabaseTransaction,
        user_id: String,
    ) -> Result<Vec<SignaturePolicy>, PolicyError> {
        let policy_models = PolicyEntity::find()
            .filter(Column::UserId.eq(user_id))
            .into_model::<Model>()
            .all(db)
            .await
            .map_err(|e| PolicyError::DatabaseOperationError(e.to_string()))?;

        let policies: Vec<SignaturePolicy> = policy_models
            .into_iter()
            .map(|model| SignaturePolicy::from(model))
            .collect();

        Ok(policies)
    }

    /// Marks a policy as corrupted by setting its valid code to 1.
    ///
    /// # Arguments
    /// * `db` - The database transaction
    /// * `policy_id` - The UUID of the policy to mark as corrupted
    ///
    /// # Returns
    /// * `Result<(), PolicyError>` - Ok if successful, otherwise a PolicyError
    pub async fn set_is_corrupted(
        db: &DatabaseTransaction,
        policy_id: Uuid,
    ) -> Result<(), PolicyError> {
        let policy = PolicyEntity::find_by_id(policy_id.to_string())
            .one(db)
            .await
            .map_err(|e| PolicyError::DatabaseOperationError(e.to_string()))?
            .ok_or_else(|| PolicyError::PolicyNotFoundError("Policy not found".to_string()))?;
        let mut active_model: ActiveModel = policy.into();
        active_model.valid_code = ActiveValue::Set(1);
        active_model
            .update(db)
            .await
            .map_err(|e| PolicyError::DatabaseOperationError(e.to_string()))?;
        Ok(())
    }

    /// Marks a policy as corrupted using a database connection instead of a transaction.
    ///
    /// # Arguments
    /// * `db` - The database connection
    /// * `policy_id` - The UUID of the policy to mark as corrupted
    ///
    /// # Returns
    /// * `Result<(), PolicyError>` - Ok if successful, otherwise a PolicyError
    pub async fn set_is_corrupted_use_connection(
        db: &DatabaseConnection,
        policy_id: Uuid,
    ) -> Result<(), PolicyError> {
        let policy = PolicyEntity::find_by_id(policy_id.to_string())
            .one(db)
            .await
            .map_err(|e| PolicyError::DatabaseOperationError(e.to_string()))?
            .ok_or_else(|| PolicyError::PolicyNotFoundError("Policy not found".to_string()))?;
        let mut active_model: ActiveModel = policy.into();
        active_model.valid_code = ActiveValue::Set(1);
        active_model
            .update(db)
            .await
            .map_err(|e| PolicyError::DatabaseOperationError(e.to_string()))?;
        Ok(())
    }

    /// Retrieves non-corrupted policies by their IDs.
    ///
    /// # Arguments
    /// * `db` - The database connection
    /// * `uuids` - List of policy UUIDs to retrieve
    ///
    /// # Returns
    /// * `Result<Vec<SignaturePolicy>, PolicyError>` - List of non-corrupted policies if successful, otherwise a PolicyError
    pub async fn get_correct_policies_by_ids(
        db: &DatabaseConnection,
        uuids: Vec<Uuid>,
    ) -> Result<Vec<SignaturePolicy>, PolicyError> {
        let policy_models = PolicyEntity::find()
            .filter(Column::PolicyId.is_in(uuids).and(Column::ValidCode.eq(0)))
            .into_model::<Model>()
            .all(db)
            .await
            .map_err(|e| PolicyError::DatabaseOperationError(e.to_string()))?;
        let signature_policy = policy_models
            .into_iter()
            .map(|model| model.into())
            .collect::<Vec<_>>();
        Ok(signature_policy)
    }

    /// Gets the total count of policies owned by a specific user.
    ///
    /// # Arguments
    /// * `db` - The database transaction
    /// * `user_id` - The ID of the user whose policies to count
    ///
    /// # Returns
    /// * `Result<u64, PolicyError>` - The count of policies if successful, otherwise a PolicyError
    pub async fn get_user_policy_count(
        db: &DatabaseTransaction,
        user_id: String,
    ) -> Result<u64, PolicyError> {
        PolicyEntity::find()
            .filter(Column::UserId.eq(user_id))
            .count(db)
            .await
            .map_err(|e| PolicyError::DatabaseOperationError(e.to_string()))
    }

    /// Retrieves all policies that don't match the specified key version.
    ///
    /// # Arguments
    /// * `db` - The database transaction
    /// * `key_version` - The key version to compare against
    ///
    /// # Returns
    /// * `Result<Vec<SignaturePolicy>, KeyManagerError>` - List of policies with different key versions if successful, otherwise a KeyManagerError
    pub async fn get_policies_by_key_version(
        db: &DatabaseTransaction,
        key_version: &str,
    ) -> Result<Vec<SignaturePolicy>, KeyManagerError> {
        let result = PolicyEntity::find()
            .filter(Column::KeyVersion.ne(key_version))
            .into_model::<Model>()
            .all(db)
            .await
            .map_err(|e| KeyManagerError::new(format!("Database operation error: {}", e)))?;

        let policies = result.into_iter().map(SignaturePolicy::from).collect();
        Ok(policies)
    }

    /// Updates the signature and key version of a policy.
    ///
    /// # Arguments
    /// * `db` - The database transaction
    /// * `policy_id` - The ID of the policy to update
    /// * `new_version` - The new key version
    /// * `new_signature` - The new signature bytes
    ///
    /// # Returns
    /// * `Result<(), KeyManagerError>` - Ok if successful, otherwise a KeyManagerError
    pub async fn update_policy_signature(
        db: &DatabaseTransaction,
        policy_id: String,
        new_version: &str,
        new_signature: &[u8],
    ) -> Result<(), KeyManagerError> {
        let policy = PolicyEntity::find_by_id(policy_id)
            .one(db)
            .await
            .map_err(|e| KeyManagerError::new(format!("Failed to find policy: {}", e)))?;

        if let Some(model) = policy {
            let mut active_model: ActiveModel = model.into();
            active_model.key_version = Set(new_version.to_string());
            active_model.signature = Set(new_signature.to_vec());

            active_model.update(db).await.map_err(|e| {
                KeyManagerError::new(format!("Failed to update policy signature: {}", e))
            })?;
        }

        Ok(())
    }

    /// Updates the corruption status of a policy.
    ///
    /// # Arguments
    /// * `db` - The database transaction
    /// * `policy_id` - The ID of the policy to update
    /// * `vaild_code` - The new validity code to set
    ///
    /// # Returns
    /// * `Result<(), KeyManagerError>` - Ok if successful, otherwise a KeyManagerError
    pub async fn update_policy_corrupted(
        db: &DatabaseTransaction,
        policy_id: String,
        vaild_code: i8,
    ) -> Result<(), KeyManagerError> {
        let policy = PolicyEntity::find_by_id(policy_id)
            .one(db)
            .await
            .map_err(|e| KeyManagerError::new(format!("Failed to find policy: {}", e)))?;

        if let Some(model) = policy {
            let mut active_model: ActiveModel = model.into();
            active_model.valid_code = Set(vaild_code);

            active_model.update(db).await.map_err(|e| {
                KeyManagerError::new(format!("Failed to update policy corrupted status: {}", e))
            })?;
        }

        Ok(())
    }
}
